import pandas as pd
import numpy as np


# 假设df是包含股票数据的DataFrame，且包含'close'列
def calculate_macd(df):
    # 计算DIF
    df['DIF'] = df['close'].ewm(span=9, adjust=False).mean() - df['close'].ewm(span=150, adjust=False).mean()

    # 计算DEA
    df['DEA'] = df['DIF'].ewm(span=12, adjust=False).mean()

    # 计算MACD
    df['MACD'] = 2 * (df['DIF'] - df['DEA'])

    return df


# 检测启动点
def macd_start_point(df):
    # 计算启动点
    conditions = [
        df['MACD'] < 0,
        df['MACD'] > df['MACD'].shift(1),
        df['MACD'].shift(1) < df['MACD'].shift(2),
        df['MACD'].shift(2) < df['MACD'].shift(3),
        df['MACD'].shift(3) < df['MACD'].shift(4)
    ]
    df['启动点'] = np.where(np.all(conditions, axis=0), 1, 0)

    # # 近2日和今日的启动点
    # df['近2日'] = np.where(df['启动点'] | df['启动点'].shift(1), 1, 0)
    # df['今日'] = np.where(df['启动点'], 1, 0)

    return df

if __name__ == '__main__':
    df = pd.read_csv('../data_info/bit_coin.csv')
    # 应用MACD和启动点检测
    df = calculate_macd(df)
    df = macd_start_point(df)
    print(df.head())
